import os
import json
import matplotlib.pyplot as plt
import seaborn as sns
from config import ROOT_DIR
import numpy as np

def load_metrics(metrics_file):
    """加载训练指标文件"""
    with open(metrics_file, 'r', encoding='utf-8') as f:
        return json.load(f)

def plot_training_curves(metrics, save_dir):
    """绘制训练曲线"""
    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    
    # 创建保存目录
    os.makedirs(save_dir, exist_ok=True)
    
    # 设置绘图风格
    plt.style.use('seaborn')
    
    # 1. 损失曲线
    plt.figure(figsize=(10, 6))
    plt.plot(metrics['epochs'], metrics['train_loss'], label='训练损失', linewidth=2)
    plt.plot(metrics['epochs'], metrics['val_loss'], label='验证损失', linewidth=2)
    plt.title('训练和验证损失曲线')
    plt.xlabel('Epoch')
    plt.ylabel('损失')
    plt.legend()
    plt.grid(True)
    plt.savefig(os.path.join(save_dir, 'loss_curves.png'))
    plt.close()
    
    # 2. AUC曲线
    plt.figure(figsize=(10, 6))
    train_auc = [m['auc'] for m in metrics['train_metrics']]
    val_auc = [m['auc'] for m in metrics['val_metrics']]
    plt.plot(metrics['epochs'], train_auc, label='训练AUC', linewidth=2)
    plt.plot(metrics['epochs'], val_auc, label='验证AUC', linewidth=2)
    plt.title('训练和验证AUC曲线')
    plt.xlabel('Epoch')
    plt.ylabel('AUC')
    plt.legend()
    plt.grid(True)
    plt.savefig(os.path.join(save_dir, 'auc_curves.png'))
    plt.close()
    
    # 3. 精确率、召回率、F1分数曲线
    plt.figure(figsize=(12, 6))
    train_precision = [m['precision'] for m in metrics['train_metrics']]
    train_recall = [m['recall'] for m in metrics['train_metrics']]
    train_f1 = [m['f1'] for m in metrics['train_metrics']]
    
    plt.plot(metrics['epochs'], train_precision, label='训练精确率', linewidth=2)
    plt.plot(metrics['epochs'], train_recall, label='训练召回率', linewidth=2)
    plt.plot(metrics['epochs'], train_f1, label='训练F1分数', linewidth=2)
    
    plt.title('训练集评估指标曲线')
    plt.xlabel('Epoch')
    plt.ylabel('分数')
    plt.legend()
    plt.grid(True)
    plt.savefig(os.path.join(save_dir, 'train_metrics_curves.png'))
    plt.close()
    
    # 4. 验证集评估指标曲线
    plt.figure(figsize=(12, 6))
    val_precision = [m['precision'] for m in metrics['val_metrics']]
    val_recall = [m['recall'] for m in metrics['val_metrics']]
    val_f1 = [m['f1'] for m in metrics['val_metrics']]
    
    plt.plot(metrics['epochs'], val_precision, label='验证精确率', linewidth=2)
    plt.plot(metrics['epochs'], val_recall, label='验证召回率', linewidth=2)
    plt.plot(metrics['epochs'], val_f1, label='验证F1分数', linewidth=2)
    
    plt.title('验证集评估指标曲线')
    plt.xlabel('Epoch')
    plt.ylabel('分数')
    plt.legend()
    plt.grid(True)
    plt.savefig(os.path.join(save_dir, 'val_metrics_curves.png'))
    plt.close()
    
    # 5. 混淆矩阵热力图（最佳epoch）
    best_epoch_idx = metrics['best_epoch'] - 1
    best_val_conf_matrix = np.array(metrics['val_metrics'][best_epoch_idx]['confusion_matrix'])
    
    plt.figure(figsize=(8, 6))
    sns.heatmap(best_val_conf_matrix, annot=True, fmt='d', cmap='Blues')
    plt.title(f'最佳验证集混淆矩阵 (Epoch {metrics["best_epoch"]})')
    plt.xlabel('预测标签')
    plt.ylabel('真实标签')
    plt.savefig(os.path.join(save_dir, 'best_confusion_matrix.png'))
    plt.close()
    
    # 6. 生成训练报告
    report = f"""
训练报告
========
模型信息:
- 模型名称: {metrics['model_info']['model_name']}
- 训练设备: {metrics['model_info']['device']}
- 优化器: {metrics['model_info']['optimizer']}
- 学习率: {metrics['model_info']['learning_rate']}
- 训练轮数: {metrics['model_info']['num_epochs']}
- 批次大小: {metrics['model_info']['batch_size']}
- 训练集大小: {metrics['model_info']['train_size']}
- 验证集大小: {metrics['model_info']['val_size']}

训练结果:
- 最佳验证AUC: {metrics['best_val_auc']:.4f}
- 最佳Epoch: {metrics['best_epoch']}
- 总训练时间: {metrics['training_time']:.2f}秒
- 平均每轮时间: {metrics['training_time']/metrics['model_info']['num_epochs']:.2f}秒

最终评估指标:
训练集:
- AUC: {metrics['train_metrics'][-1]['auc']:.4f}
- 精确率: {metrics['train_metrics'][-1]['precision']:.4f}
- 召回率: {metrics['train_metrics'][-1]['recall']:.4f}
- F1分数: {metrics['train_metrics'][-1]['f1']:.4f}
- 特异度: {metrics['train_metrics'][-1]['specificity']:.4f}

验证集:
- AUC: {metrics['val_metrics'][-1]['auc']:.4f}
- 精确率: {metrics['val_metrics'][-1]['precision']:.4f}
- 召回率: {metrics['val_metrics'][-1]['recall']:.4f}
- F1分数: {metrics['val_metrics'][-1]['f1']:.4f}
- 特异度: {metrics['val_metrics'][-1]['specificity']:.4f}
"""
    
    # 保存训练报告
    with open(os.path.join(save_dir, 'training_report.txt'), 'w', encoding='utf-8') as f:
        f.write(report)

def main():
    # 获取最新的训练指标文件
    metrics_dir = os.path.join(ROOT_DIR, 'training_metrics')
    metrics_files = [f for f in os.listdir(metrics_dir) if f.endswith('.json')]
    if not metrics_files:
        print("未找到训练指标文件！")
        return
    
    latest_metrics_file = max(metrics_files, key=lambda x: os.path.getctime(os.path.join(metrics_dir, x)))
    metrics_file_path = os.path.join(metrics_dir, latest_metrics_file)
    
    # 加载指标
    metrics = load_metrics(metrics_file_path)
    
    # 创建图表保存目录
    plots_dir = os.path.join(ROOT_DIR, 'training_plots')
    
    # 绘制训练曲线
    plot_training_curves(metrics, plots_dir)
    print(f"训练曲线已保存到: {plots_dir}")

if __name__ == "__main__":
    main() 